import os

import cv2
import numpy as np
import torchvision
from torch.utils.tensorboard import SummaryWriter


# Writer will output to ./runs/ directory by default

def cv_show_image_single(image, dtype=np.uint8):
    image = np.array(image, dtype=np.uint8)
    cv2.destroyAllWindows()
    win_caption = "ImageViewer"
    cv2.namedWindow(win_caption)
    cv2.imshow(win_caption, image)
    cv2.waitKey()


class TensorboardVisualizer(object):

    def __init__(self, log_dir='./logs/'):
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
        self.writer = SummaryWriter(log_dir=log_dir)

    def add_image_batch(self, image_batch, n_iter, image_name='Image_batch'):
        grid = torchvision.utils.make_grid(image_batch)
        self.writer.add_image(image_name, grid, n_iter)

    def add_image_one(self, image, n_iter, image_name='Image_single'):
        grid = torchvision.utils.make_grid(image)
        self.writer.add_image(image_name, grid, n_iter)

    def add_loss_step(self, loss_val, n_iter, loss_name='loss'):
        self.writer.add_scalar(loss_name, loss_val, n_iter)

    def add_graph_image(self, model_net, input_batch_images):
        self.writer.add_graph(model_net, input_batch_images)

    def add_graph_shape(self, model_net, input_shape):
        input_image = np.zeros(input_shape)
        self.writer.add_graph(model_net, input_image)

    def close(self):
        self.writer.close()
